package top.continew.ai.config;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;

/**
 * 查询重写器
 *
 * @Author： Lonni
 * @Date： 2025/8/21  19:24
 */
public class QueryRewriter {

    public final QueryTransformer queryTransformer;


    public QueryRewriter(ChatClient.Builder builder) {
        RewriteQueryTransformer build = RewriteQueryTransformer.builder()
                .chatClientBuilder( builder)
                .build();

        this.queryTransformer = build;
    }

    public String rewrite(String prompt) {
        Query query = new Query(prompt);
        // 执行查询重写
        Query transformedQuery = queryTransformer.transform(query);
        // 输出重写后的查询
        return transformedQuery.text();
    }


}
